{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RGGFAk4ok3W2"
      },
      "source": [
        "Copyright 2021 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "     https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfnC1bgpk43k"
      },
      "source": [
        "# Finite ConvNet Training with Conv-KIP distilled images\n",
        "\n",
        "This notebook demonstrate a simple example of finite neural network transfer using Conv-KIP distilled image. \n",
        "\n",
        "Training code is based off of [Lee et al., Finite Versus Infinite Neural Networks: an Empirical Study, NeurIPS 2020](https://arxiv.org/abs/2007.15801), as adapted in [Nguyen et al., Dataset Distillation with Infinitely Wide Convolutional Networks](https://arxiv.org/abs/2107.13034)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tpXmXTFgpNKC"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_jZmTD5hpDkG"
      },
      "outputs": [],
      "source": [
        "# Install ml_collections and neural_tangents\n",
        "!pip install -q git+https://www.github.com/google/ml_collections\n",
        "!pip install -q git+https://www.github.com/google/neural-tangents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HCN_G3bA7WU_"
      },
      "outputs": [],
      "source": [
        "from absl import app\n",
        "from absl import logging\n",
        "\n",
        "import functools\n",
        "import time\n",
        "import operator as op\n",
        "\n",
        "import jax\n",
        "from jax.example_libraries import optimizers\n",
        "from jax.example_libraries import stax as ostax\n",
        "import jax.numpy as jnp\n",
        "from jax.tree_util import tree_map\n",
        "from jax.tree_util import tree_reduce\n",
        "\n",
        "import ml_collections\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import stax\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iaWIPFZvmtEE"
      },
      "source": [
        "## Experiment Configs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uIz0K4QEmwpU"
      },
      "outputs": [],
      "source": [
        "def get_config():\n",
        "  # Note that max_lr_factor and l2_regularization is found through grid search.\n",
        "  config = ml_collections.ConfigDict()\n",
        "  config.seed = 0\n",
        "\n",
        "  # Dataset\n",
        "  config.dataset = ml_collections.ConfigDict()\n",
        "  config.dataset.name = 'cifar10'  # ['cifar10', 'cifar100', 'mnist', 'fashion_mnist']\n",
        "\n",
        "  config.preprocess_type = 'zca_normalize'  # ['zca_normalize', 'standard']\n",
        "  config.zca_reg = 0.1\n",
        "\n",
        "  # Optimization\n",
        "  config.optimizer = 'momentum' # Note: In this notebook `run_trainer` optimizer is hard-coded to be `momentum`.\n",
        "  config.momentum = 0.9\n",
        "  config.batch_size = 100\n",
        "  config.eval_batch_size = 100\n",
        "  config.train_steps = 5_000\n",
        "\n",
        "  config.empirical_ntk_num_inputs = 50  # Number of samples to estimate max LR.\n",
        "\n",
        "  config.max_lr_factor = 1.0\n",
        "  config.l2_regularization = 0.001\n",
        "\n",
        "  # Network Architecture\n",
        "  config.width = 1024\n",
        "  config.depth = 3\n",
        "  config.use_gap = False\n",
        "  config.W_std = 2.0**0.5\n",
        "  config.b_std = 0.1\n",
        "  config.activation = 'relu'  # ['relu', 'identity', 'gelu'', 'identity']\n",
        "  config.prepend_conv_layer = True  # enable for 'ConvNet' datasets; disable for 'ConvNet3' datasets\n",
        "\n",
        "  config.loss_type = 'mse'  # ['mse', 'xent']\n",
        "  config.parameterization = 'standard'  # ['standard', 'ntk']\n",
        "\n",
        "  config.kip = ml_collections.ConfigDict()\n",
        "  # Put any KIP / Label Solve checkpoint path here to use as training data.\"\n",
        "  config.kip.data_ckpt_path = (\n",
        "      'gs://kip-datasets/kip/cifar10/ConvNet_ssize100_zca_nol_noaug_ckpt1000.npz')\n",
        "\n",
        "  return config"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BwO9GQYkmWfI"
      },
      "source": [
        "## Define Training Utilities"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SLBLnZz47uoW"
      },
      "outputs": [],
      "source": [
        "def load_kip_data(config):\n",
        "  \"\"\"Load and preprocess dataset with TFDS and KIP ckpt path.\"\"\"\n",
        "  data = _load_dataset(config)\n",
        "  data = apply_preprocess(config, data)\n",
        "\n",
        "  logging.info('valid data: %s, %s', \n",
        "               data['valid']['images'].shape, data['valid']['labels'].shape)\n",
        "  logging.info('test data: %s, %s', \n",
        "               data['test']['images'].shape, data['test']['labels'].shape)\n",
        "\n",
        "  # Override training data\n",
        "  ckpt_name = config.kip.data_ckpt_path\n",
        "  with tf.io.gfile.GFile(ckpt_name, 'rb') as f:\n",
        "    loaded_data = jnp.load(f)\n",
        "  train_images = loaded_data['images']\n",
        "  train_labels = loaded_data['labels']\n",
        "\n",
        "  if config.loss_type == 'xent':\n",
        "    # Recover one-hot label even for the learned labels.\n",
        "    n_classes = train_labels.shape[-1]\n",
        "    train_labels = jnp.array(\n",
        "        jnp.argmax(train_labels, axis=-1)[:, None] == jnp.arange(n_classes),\n",
        "        dtype=train_labels.dtype)\n",
        "\n",
        "  # Override training\n",
        "  data['train'] = {'images': train_images, 'labels': train_labels}\n",
        "\n",
        "  logging.info('Overrriding train with ckpt %s, size: (%s, %s)', ckpt_name,\n",
        "               data['train']['images'].shape, data['train']['labels'].shape)\n",
        "\n",
        "  return data\n",
        "\n",
        "\n",
        "def _load_dataset(config):\n",
        "  \"\"\"Get per channel normalized / one hot encoded data from TFDS.\"\"\"\n",
        "  VALID_SIZE = 5000\n",
        "\n",
        "  dataset_name = config.dataset.name\n",
        "  ds_builder = tfds.builder(dataset_name)\n",
        "  ds_train, ds_test = tfds.as_numpy(\n",
        "      tfds.load(\n",
        "          dataset_name,\n",
        "          split=['train', 'test'],\n",
        "          batch_size=-1,\n",
        "          as_dataset_kwargs={'shuffle_files': False}))\n",
        "\n",
        "  train_images, train_labels, test_images, test_labels = (ds_train['image'],\n",
        "                                                          ds_train['label'],\n",
        "                                                          ds_test['image'],\n",
        "                                                          ds_test['label'])\n",
        "  height, width, num_channels = ds_builder.info.features['image'].shape\n",
        "  num_classes = ds_builder.info.features['label'].num_classes\n",
        "  with config.dataset.unlocked():\n",
        "    config.dataset.height = height\n",
        "    config.dataset.width = width\n",
        "    config.dataset.num_channels = num_channels\n",
        "    config.dataset.num_classes = num_classes\n",
        "\n",
        "  # One hot encode\n",
        "  train_labels = jax.nn.one_hot(train_labels, num_classes)\n",
        "  test_labels = jax.nn.one_hot(test_labels, num_classes)\n",
        "\n",
        "  if config.get('loss_type', 'mse') == 'mse':\n",
        "    shift = (1. / num_classes if num_classes \u003e 1 else 0.5)\n",
        "    train_labels -= shift\n",
        "    test_labels -= shift\n",
        "\n",
        "  # Normalize by precomputed per channel mean/std from training images\n",
        "  train_xs = (train_images - np.mean(train_images, axis=(0, 1, 2))) / np.std(\n",
        "      train_images, axis=(0, 1, 2))\n",
        "  test_xs = (test_images - np.mean(train_images, axis=(0, 1, 2))) / np.std(\n",
        "      train_images, axis=(0, 1, 2))\n",
        "\n",
        "  test_ys = test_labels\n",
        "  train_xs, valid_xs = train_xs[:-VALID_SIZE], train_xs[-VALID_SIZE:]\n",
        "  train_ys, valid_ys = train_labels[:-VALID_SIZE], train_labels[-VALID_SIZE:]\n",
        "\n",
        "  train = (train_xs, train_ys)\n",
        "  valid = (valid_xs, valid_ys)\n",
        "  test = (test_xs, test_ys)\n",
        "  \n",
        "  data = {'train': {'images': train_xs, 'labels': train_ys},\n",
        "          'valid': {'images': valid_xs, 'labels': valid_ys},\n",
        "          'test': {'images': test_xs, 'labels': test_ys}\n",
        "  }\n",
        "\n",
        "  return data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "eUwrmFWxfuAf"
      },
      "outputs": [],
      "source": [
        "#@title Preprocess utilties\n",
        "def apply_preprocess(config, data):\n",
        "  \"\"\"Apply ZCA preprocessing on the standard normalized data.\"\"\"\n",
        "  x_train, y_train = data['train']['images'], data['train']['labels']\n",
        "  x_valid, y_valid = data['valid']['images'], data['valid']['labels']\n",
        "  x_test, y_test = data['test']['images'], data['test']['labels']\n",
        "\n",
        "  preprocess_type = config.get('preprocess_type', 'standard')\n",
        "  if preprocess_type == 'standard':\n",
        "    # Normalization is already done.\n",
        "    pass\n",
        "  else:\n",
        "    zca_reg = config.get('zca_reg', 0.0)\n",
        "    if preprocess_type == 'zca_normalize':\n",
        "        preprocess_op = _get_preprocess_op(\n",
        "          x_train,\n",
        "          layer_norm=True,\n",
        "          zca_reg=zca_reg,\n",
        "          zca_reg_absolute_scale=config.get('zca_reg_absolute_scale', False))\n",
        "        x_train = preprocess_op(x_train)\n",
        "        x_valid = preprocess_op(x_valid)\n",
        "        x_test = preprocess_op(x_test)\n",
        "    else:\n",
        "      NotImplementedError('Preprocess type %s is not implemented' %\n",
        "                          preprocess_type)\n",
        "\n",
        "  return {'train': {'images': x_train, 'labels': y_train},\n",
        "          'valid': {'images': x_valid, 'labels': y_valid},\n",
        "          'test': {'images': x_test, 'labels': y_test}}\n",
        "\n",
        "\n",
        "def _get_preprocess_op(x_train,\n",
        "                      layer_norm=True,\n",
        "                      zca_reg=1e-5,\n",
        "                      zca_reg_absolute_scale=False,\n",
        "                      on_cpu=False):\n",
        "  \"\"\"ZCA preprocessing function.\"\"\"\n",
        "  whitening_transform = _get_whitening_transform(x_train, layer_norm, zca_reg,\n",
        "                                                zca_reg_absolute_scale,\n",
        "                                                on_cpu)\n",
        "\n",
        "  def _preprocess_op(images):\n",
        "    orig_shape = images.shape\n",
        "    images = images.reshape(orig_shape[0], -1)\n",
        "    if layer_norm:\n",
        "      # Zero mean every feature\n",
        "      images = images - jnp.mean(images, axis=1)[:, jnp.newaxis]\n",
        "      # Normalize\n",
        "      image_norms = jnp.linalg.norm(images, axis=1)\n",
        "      # Make features unit norm\n",
        "      images = images / image_norms[:, jnp.newaxis]\n",
        "\n",
        "    images = (images).dot(whitening_transform)\n",
        "    images = images.reshape(orig_shape)\n",
        "    return images\n",
        "\n",
        "  return _preprocess_op\n",
        "\n",
        "\n",
        "def _get_whitening_transform(x_train,\n",
        "                             layer_norm=True,\n",
        "                             zca_reg=1e-5,\n",
        "                             zca_reg_absolute_scale=False,\n",
        "                             on_cpu=False):\n",
        "  \"\"\"Returns 2D matrix that performs whitening transform.\n",
        "\n",
        "  Whitening transform is a (d,d) matrix (d = number of features) which acts on\n",
        "  the right of a (n, d) batch of flattened data.\n",
        "  \"\"\"\n",
        "  orig_train_shape = x_train.shape\n",
        "  x_train = x_train.reshape(orig_train_shape[0], -1).astype('float64')\n",
        "  if on_cpu:\n",
        "    x_train = jax.device_put(x_train, jax.devices('cpu')[0])\n",
        "\n",
        "  n_train = x_train.shape[0]\n",
        "  if layer_norm:\n",
        "    logging.info('Performing layer norm preprocessing.')\n",
        "    # Zero mean every feature\n",
        "    x_train = x_train - jnp.mean(x_train, axis=1)[:, jnp.newaxis]\n",
        "    # Normalize\n",
        "    train_norms = jnp.linalg.norm(x_train, axis=1)\n",
        "    # Make features unit norm\n",
        "    x_train = x_train / train_norms[:, jnp.newaxis]\n",
        "\n",
        "  logging.info('Performing zca whitening preprocessing with reg: %.2e', zca_reg)\n",
        "  cov = 1.0 / n_train * x_train.T.dot(x_train)\n",
        "  if zca_reg_absolute_scale:\n",
        "    reg_amount = zca_reg\n",
        "  else:\n",
        "    reg_amount = zca_reg * jnp.trace(cov) / cov.shape[0]\n",
        "  logging.info('Raw zca regularization strength: %f', reg_amount)\n",
        "\n",
        "  u, s, _ = jnp.linalg.svd(cov + reg_amount * jnp.eye(cov.shape[0]))\n",
        "  inv_sqrt_zca_eigs = s**(-1 / 2)\n",
        "\n",
        "  # rank control\n",
        "  if n_train \u003c x_train.shape[1]:\n",
        "    inv_sqrt_zca_eigs = inv_sqrt_zca_eigs.at[n_train:].set(\n",
        "        jnp.ones(inv_sqrt_zca_eigs[n_train:].shape[0]))\n",
        "  whitening_transform = jnp.einsum(\n",
        "      'ij,j,kj-\u003eik', u, inv_sqrt_zca_eigs, u, optimize=True)\n",
        "  return whitening_transform\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "StI0M1sd9pN0"
      },
      "outputs": [],
      "source": [
        "# Loss Definition.\n",
        "cross_entropy = lambda y, y_hat: -np.mean(np.sum(y * y_hat, axis=1))\n",
        "mse_loss = lambda y, y_hat: 0.5 * jnp.mean((y - y_hat)**2)\n",
        "\n",
        "_l2_norm = lambda params: tree_map(lambda x: jnp.sum(x ** 2), params)\n",
        "l2_regularization = lambda params: tree_reduce(op.add, _l2_norm(params))\n",
        "\n",
        "\n",
        "def cosine_schedule(initial_learning_rate, training_steps):\n",
        "  def _cosine_schedule(t):\n",
        "    return initial_learning_rate * 0.5 * (\n",
        "        1 + jnp.cos(t / training_steps * jnp.pi))\n",
        "  return _cosine_schedule\n",
        "\n",
        "def _epoch_from_step(step, train_size, batch_size):\n",
        "  if train_size == batch_size:\n",
        "    return step\n",
        "  else:\n",
        "    return float(step / train_size * batch_size)  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B2NJelTBpBL_"
      },
      "source": [
        "## Define Networks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ehZ4P7BpCv2"
      },
      "outputs": [],
      "source": [
        "def _get_activation_fn(config):\n",
        "  if config.activation.lower() == 'relu':\n",
        "    activation_fn = stax.Relu()\n",
        "  elif config.activation.lower() == 'erf':\n",
        "    activation_fn = stax.Erf()\n",
        "  elif config.activation.lower() == 'identity':\n",
        "    activation_fn = stax.Identity()\n",
        "  elif config.activation.lower() == 'gelu':\n",
        "    activation_fn = stax.Gelu()\n",
        "  else:\n",
        "    raise ValueError('activation function %s not implemented' %\n",
        "                     config.activation)\n",
        "  return activation_fn\n",
        "\n",
        "\n",
        "def _get_norm_layer(normalization):\n",
        "  normalization = normalization.lower()\n",
        "  if 'layer' in normalization:\n",
        "    norm_layer = stax.LayerNorm(axis=(1, 2, 3))\n",
        "  elif 'instance' in normalization:\n",
        "    norm_layer = stax.LayerNorm(axis=(1, 2))\n",
        "  elif normalization == '':\n",
        "    norm_layer = stax.Identity()\n",
        "  else:\n",
        "    raise ValueError('normalization %s not implemented' % normalization)\n",
        "  return norm_layer\n",
        " \n",
        "\n",
        "def _ConvNet(config): \n",
        "  return ConvNet(\n",
        "      depth=config.depth,\n",
        "      width=config.width,\n",
        "      prepend_conv_layer=config.prepend_conv_layer,\n",
        "      use_gap=config.get('use_gap', False),\n",
        "      W_std=config.W_std,\n",
        "      b_std=config.b_std,\n",
        "      num_classes=config.dataset.num_classes,\n",
        "      parameterization=config.parameterization,\n",
        "      activation_fn=_get_activation_fn(config),\n",
        "      norm_layer=_get_norm_layer(config.get('normalization', '')),\n",
        "      image_format=config.get('image_format', 'NHWC'))\n",
        "\n",
        "\n",
        "def ConvNet( \n",
        "    depth: int,\n",
        "    width: int,\n",
        "    prepend_conv_layer: bool = True,\n",
        "    use_gap: bool = False,\n",
        "    W_std=2**0.5,\n",
        "    b_std=0.1,\n",
        "    num_classes: int = 10,\n",
        "    parameterization: str = 'ntk',\n",
        "    activation_fn=stax.Relu(),\n",
        "    norm_layer=stax.Identity(),\n",
        "    image_format: str = 'NHWC'):\n",
        "  \"\"\"Adaptation of ConvNet baseline of Dataset Condensation.\n",
        "\n",
        "  Original architecture is based on (Gidaris \u0026 Komodakis, 2018)\n",
        "  and here we adapt version of Zhao et al., Dataset Condensation with Gradient\n",
        "  Matching, https://openreview.net/pdf?id=mSAKhLYLSsl\n",
        "\n",
        "  Implements depth-many blocks of convolution, activation, 2x2 avg pooling.\n",
        "  Normalization layer of corresponding finite-width neural network is omitted.\n",
        "\n",
        "  For the 'ConvNet' settings of Nguyen et al., Dataset Distillation with \n",
        "  Infinitely Wide Convolutional Networks, set depth=3 (width is immaterial for\n",
        "  'ntk' parameterization) and prepend_conv_layer=True. For 'ConvNet3' settings\n",
        "  closer to Zhao et al., set prepend_conv_layer=False. \n",
        "\n",
        "  Args:\n",
        "    depth: depth of network\n",
        "    width: width of network\n",
        "    prepend_conv_layer: if True, add an additional conv and relu layer before\n",
        "      main set of blocks\n",
        "    use_gap: if True, use global average pooling for preclassifier layer\n",
        "    W_std: standard deviation of weight matrix initialization\n",
        "    b_std: standard deviation of bias initialization\n",
        "    num_classes: number of classes for output layer\n",
        "    parameterization: 'ntk' or 'standard' for initializing network and NTK\n",
        "    activation_fn: NT activation function of network\n",
        "    norm_layer: NT normalization layer, default is Identity.\n",
        "    image_format: Image format 'NHWC', 'NCHW' etc.\n",
        "\n",
        "  Returns:\n",
        "    Corresponding neural_tangents stax model.\n",
        "  \"\"\"\n",
        "  layers = []\n",
        "  conv = functools.partial(\n",
        "      stax.Conv,\n",
        "      W_std=W_std,\n",
        "      b_std=b_std,\n",
        "      padding='SAME',\n",
        "      parameterization=parameterization)\n",
        "  if prepend_conv_layer:\n",
        "    layers += [\n",
        "        conv(width, (3, 3),\n",
        "             dimension_numbers=(image_format, 'HWIO', 'NHWC')),\n",
        "        activation_fn\n",
        "    ]\n",
        "\n",
        "  # generate blocks of convolutions followed by average pooling\n",
        "  for _ in range(depth):\n",
        "    layers += [conv(width, (3, 3)), norm_layer,\n",
        "               activation_fn, stax.AvgPool((2, 2), strides=(2, 2))]\n",
        "  if use_gap:\n",
        "    layers.append(stax.GlobalAvgPool())\n",
        "  else:\n",
        "    layers.append(stax.Flatten())\n",
        "  layers.append(stax.Dense(num_classes, W_std, b_std,\n",
        "                           parameterization=parameterization))\n",
        "\n",
        "  return stax.serial(*layers)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iEd6CeN5npqo"
      },
      "source": [
        "## Define Trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nYEqYzzEkmyB"
      },
      "outputs": [],
      "source": [
        "def run_trainer(data, config):\n",
        "  \"\"\"Train a neural network.\"\"\"\n",
        "\n",
        "  # Experiment Parameters.\n",
        "  batch_size = config.batch_size\n",
        "  eval_batch_size = config.get('eval_batch_size', config.batch_size)\n",
        "  train_size = data['train']['images'].shape[0]\n",
        "  steps_per_epoch = int(np.ceil(train_size / batch_size))\n",
        "  train_steps = int(config.train_steps)\n",
        "  train_epochs = int(np.ceil(config.train_steps / steps_per_epoch))\n",
        "\n",
        "  l2_lambda = config.l2_regularization\n",
        "\n",
        "  key = jax.random.PRNGKey(config.seed)\n",
        "\n",
        "  # Construct tf.data\n",
        "  train_ds = tf.data.Dataset.from_tensor_slices({\n",
        "      'images': data['train']['images'],\n",
        "      'labels': data['train']['labels'],\n",
        "  }).repeat().shuffle(\n",
        "      data['train']['images'].shape[0], seed=0).batch(batch_size).as_numpy_iterator()\n",
        "\n",
        "  # This is used for computing training metrics. \n",
        "  train_eval_ds = tf.data.Dataset.from_tensor_slices({\n",
        "      'images': data['train']['images'],\n",
        "      'labels': data['train']['labels'],\n",
        "  }).batch(eval_batch_size)\n",
        "\n",
        "  valid_ds = tf.data.Dataset.from_tensor_slices({\n",
        "      'images': data['valid']['images'][:1000], # Smaller validation set size for notebook\n",
        "      'labels': data['valid']['labels'][:1000], \n",
        "  }).batch(eval_batch_size)\n",
        "\n",
        "  test_ds = tf.data.Dataset.from_tensor_slices({\n",
        "      'images': data['test']['images'],\n",
        "      'labels': data['test']['labels'],\n",
        "  }).batch(eval_batch_size)\n",
        "\n",
        "  # Initialize Network.\n",
        "  network = _ConvNet(config)\n",
        "  init_f, f, _ = network\n",
        "\n",
        "  key, split = jax.random.split(key)\n",
        "  _, init_params = init_f(split, (-1,) + data['train']['images'].shape[1:])\n",
        "\n",
        "  # Estimate maximum learning rate\n",
        "  def logit_reduced_f(params, x):\n",
        "    out = f(params, x)\n",
        "    return jnp.sum(out, axis=-1) / out.shape[-1]**(1 / 2)\n",
        "\n",
        "  input_sample = data['valid']['images'][:config.empirical_ntk_num_inputs]\n",
        "\n",
        "  empirical_kernel_fn = lambda x1, x2, params: nt.empirical_ntk_fn(\n",
        "      logit_reduced_f, trace_axes=(), vmap_axes=0, implementation=1)(x1, x2,\n",
        "                                                                     params)\n",
        "  empirical_kernel_fn = nt.batch(empirical_kernel_fn, batch_size=10)\n",
        "\n",
        "  logging.info('input_sample shape: %s', input_sample.shape)\n",
        "  max_lr_estimate_start = time.time()\n",
        "  kernel = empirical_kernel_fn(input_sample, None, init_params)\n",
        "  logging.info('kernel shape: %s', kernel.shape)\n",
        "\n",
        "  y_train_size = kernel.shape[0] * config.dataset.num_classes\n",
        "  assert y_train_size == data['valid']['labels'][:config.empirical_ntk_num_inputs].size\n",
        "  max_lr = nt.predict.max_learning_rate(\n",
        "      ntk_train_train=kernel, y_train_size=y_train_size, eps=1e-12)\n",
        "  print('Max LR estimate took: %.2fs'%(time.time() - max_lr_estimate_start))\n",
        "\n",
        "  learning_rate = float(max_lr * config.max_lr_factor)\n",
        "  print('max LR: %f, current LR: %f'%(max_lr, learning_rate))\n",
        "  \n",
        "  # Define Raw loss, Accuracy, and Optimizer.\n",
        "  @jax.jit\n",
        "  def raw_loss(params, batch):\n",
        "    \"\"\"Loss without weight decay.\"\"\"\n",
        "    images, labels = batch['images'], batch['labels']\n",
        "    loss_type = config.get('loss_type', 'xent')\n",
        "    if loss_type == 'xent':\n",
        "      return cross_entropy(ostax.logsoftmax(f(params, images)), labels)\n",
        "    elif loss_type == 'mse':\n",
        "      return mse_loss(f(params, images), labels)\n",
        "    else:\n",
        "      raise NotImplementedError('Loss type %s not implemented:' % loss_type)\n",
        "\n",
        "  @jax.jit\n",
        "  def loss(params, batch):\n",
        "    l2_loss = 0.5 * l2_lambda * l2_regularization(params)\n",
        "    return raw_loss(params, batch) + l2_loss\n",
        "\n",
        "  grad_loss = jax.jit(jax.grad(loss))\n",
        "\n",
        "  @jax.jit\n",
        "  def accuracy(params, batch):\n",
        "    images, labels = batch['images'], batch['labels']\n",
        "    return jnp.mean(\n",
        "        jnp.array(\n",
        "            jnp.argmax(f(params, images), axis=1) == jnp.argmax(labels, axis=1),\n",
        "            dtype=np.float32))\n",
        "\n",
        "  learning_rate_fn = cosine_schedule(learning_rate, config.train_steps)\n",
        "  print('Using momentum optimizer.')\n",
        "  opt_init_fn, opt_apply_fn, get_params = optimizers.momentum(\n",
        "      learning_rate_fn, config.momentum)\n",
        "\n",
        "  opt_apply_fn = jax.jit(opt_apply_fn)\n",
        "  state = opt_init_fn(init_params)\n",
        "  del init_params  # parameters obatined from optimizer state\n",
        "\n",
        "  # Define Update and Evaluate Function.\n",
        "  @jax.jit\n",
        "  def update(step, state, batch):\n",
        "    \"\"\"Training updates.\"\"\"\n",
        "    params = get_params(state)\n",
        "    new_step = step\n",
        "    dparams = grad_loss(params, batch)\n",
        "    return new_step + 1, opt_apply_fn(step, dparams, state)\n",
        "\n",
        "  def dataset_evaluate(state, dataset):\n",
        "    \"\"\"Compute loss and accuracy metrics over entire dataset.\"\"\"\n",
        "    params = get_params(state)\n",
        "    tot_metrics ={'raw_loss':0., 'loss': 0., 'correct': 0, 'count': 0}\n",
        "    for eval_batch in dataset.as_numpy_iterator():\n",
        "      eval_size = eval_batch['images'].shape[0]\n",
        "      tot_metrics['raw_loss'] += raw_loss(params, eval_batch) * eval_size\n",
        "      tot_metrics['loss'] += loss(params, eval_batch) * eval_size\n",
        "      tot_metrics['correct'] += accuracy(params, eval_batch) * eval_size\n",
        "      tot_metrics['count']  += eval_size\n",
        "    metric ={'raw_loss': tot_metrics['raw_loss'] / tot_metrics['count'], \n",
        "             'loss': tot_metrics['loss'] / tot_metrics['count'],\n",
        "             'accuracy': tot_metrics['correct']  / tot_metrics['count'] }\n",
        "\n",
        "    return metric\n",
        "\n",
        "  measurements = []\n",
        "  # Define logging steps.\n",
        "  log_max_steps = np.log10(train_steps)\n",
        "  log_steps = [0] + sorted(\n",
        "      list(set([int(10**t) for t in np.linspace(0.0, log_max_steps, 10)])))\n",
        "  start_time = time.time()\n",
        "  global_step = 0\n",
        "  step_time = 0\n",
        "\n",
        "  hparams_json = config.to_json_best_effort(indent=2)\n",
        "  print('hparams: %s', hparams_json)\n",
        "\n",
        "  print('Total training steps %d, steps_per_epoch %d' %\n",
        "        (train_steps, steps_per_epoch))\n",
        "  print('Step (Epoch)\\tLearning Rate\\tTrain Loss\\tValid Loss\\t'\n",
        "        'Train Acc\\tValid Acc\\tTime Elapsed\\tEval Time')\n",
        "  while global_step \u003c= train_steps:\n",
        "    i = int(global_step)\n",
        "    epoch = _epoch_from_step(i, train_size, batch_size)\n",
        "\n",
        "    if i in log_steps or i % 250 == 0 or i == train_steps:\n",
        "      eval_start_time = time.time()\n",
        "      train_metric = dataset_evaluate(state, train_eval_ds)\n",
        "      if not jnp.isfinite(train_metric['raw_loss']):\n",
        "        msg = 'NaN during Training! Terminating current trial.'\n",
        "        raise ValueError(msg)\n",
        "      valid_metric = dataset_evaluate(state, valid_ds)\n",
        "      \n",
        "      eval_time = time.time() - eval_start_time\n",
        "      elapsed_time = float(time.time() - start_time)\n",
        "      lr = float(learning_rate_fn(i))\n",
        "      measurements.append([\n",
        "          i, epoch, lr,\n",
        "          train_metric['loss'],\n",
        "          valid_metric['loss'],\n",
        "          train_metric['accuracy'],\n",
        "          valid_metric['accuracy'], elapsed_time\n",
        "      ])\n",
        "      print(\n",
        "          ('{:06d}\\t({:06.1f})\\t' + ('{:.6e}\\t' * 3) + ('{:.6f}\\t' * 4)).format(\n",
        "              i, epoch, lr, train_metric['loss'], valid_metric['loss'],\n",
        "              train_metric['accuracy'], valid_metric['accuracy'],\n",
        "              elapsed_time, eval_time))\n",
        "    global_step, state = update(global_step, state, next(train_ds))\n",
        "  \n",
        "  print('Training finished')\n",
        "  test_metric = dataset_evaluate(state, test_ds)\n",
        "  print('Step\\tEpoch\\tLearning Rate\\tTrain Loss\\tValid Loss\\tTest Loss\\t'\n",
        "        'Train Acc\\tValid Acc\\tTest Acc\\tTime Elapsed')\n",
        "  print(('{:06d}\\t({:06.1f})\\t' + ('{:.6e}\\t' * 4) + ('{:.6f}\\t' * 4)).format(\n",
        "      i, epoch, lr, train_metric['loss'], valid_metric['loss'], \n",
        "      test_metric['loss'], train_metric['accuracy'], \n",
        "      valid_metric['accuracy'], test_metric['accuracy'],\n",
        "      time.time() - start_time))\n",
        "  return measurements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pIPKnsUixAgs"
      },
      "source": [
        "## Run an Experiment with Trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 2038363,
          "status": "ok",
          "timestamp": 1642200416564,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "IMcqOycCo-Vu",
        "outputId": "27fd0ec7-ba4c-445e-aab2-aa6b5a8b084b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Max LR estimate took: 32.74s\n",
            "max LR: 0.006775, current LR: 0.006775\n",
            "Using momentum optimizer.\n",
            "hparams: %s {\n",
            "  \"seed\": 0,\n",
            "  \"dataset\": {\n",
            "    \"name\": \"cifar10\",\n",
            "    \"height\": 32,\n",
            "    \"width\": 32,\n",
            "    \"num_channels\": 3,\n",
            "    \"num_classes\": 10\n",
            "  },\n",
            "  \"preprocess_type\": \"zca_normalize\",\n",
            "  \"zca_reg\": 0.1,\n",
            "  \"optimizer\": \"momentum\",\n",
            "  \"momentum\": 0.9,\n",
            "  \"batch_size\": 100,\n",
            "  \"eval_batch_size\": 100,\n",
            "  \"train_steps\": 5000,\n",
            "  \"empirical_ntk_num_inputs\": 50,\n",
            "  \"max_lr_factor\": 1.0,\n",
            "  \"l2_regularization\": 0.001,\n",
            "  \"width\": 1024,\n",
            "  \"depth\": 3,\n",
            "  \"use_gap\": false,\n",
            "  \"W_std\": 1.4142135623730951,\n",
            "  \"b_std\": 0.1,\n",
            "  \"activation\": \"relu\",\n",
            "  \"loss_type\": \"mse\",\n",
            "  \"parameterization\": \"standard\",\n",
            "  \"kip\": {\n",
            "    \"data_ckpt_path\": \"gs://kip-datasets/kip/cifar10/ssize100_zca_nol_noaug_ckpt1000.npz\"\n",
            "  }\n",
            "}\n",
            "Total training steps 5000, steps_per_epoch 1\n",
            "Step (Epoch)\tLearning Rate\tTrain Loss\tValid Loss\tTrain Acc\tValid Acc\tTime Elapsed\tEval Time\n",
            "000000\t(0000.0)\t6.775185e-03\t4.233754e+00\t4.299506e+00\t0.100000\t0.117000\t8.292593\t8.292144\t\n",
            "000001\t(0001.0)\t6.775184e-03\t4.171060e+00\t4.179952e+00\t0.090000\t0.088000\t17.262976\t4.836443\t\n",
            "000002\t(0002.0)\t6.775183e-03\t4.229721e+00\t4.288811e+00\t0.090000\t0.099000\t24.235314\t4.842078\t\n",
            "000006\t(0006.0)\t6.775161e-03\t4.206002e+00\t4.247955e+00\t0.110000\t0.117000\t32.189570\t4.849226\t\n",
            "000017\t(0017.0)\t6.774992e-03\t4.153734e+00\t4.161946e+00\t0.340000\t0.133000\t42.420592\t4.860034\t\n",
            "000044\t(0044.0)\t6.773890e-03\t4.131576e+00\t4.144116e+00\t0.840000\t0.206000\t58.563524\t5.294125\t\n",
            "000113\t(0113.0)\t6.766650e-03\t4.081195e+00\t4.102843e+00\t1.000000\t0.331000\t89.776809\t5.295896\t\n",
            "000250\t(0250.0)\t6.733478e-03\t3.994498e+00\t4.024379e+00\t1.000000\t0.499000\t142.218859\t4.874434\t\n",
            "000292\t(0292.0)\t6.718330e-03\t3.970077e+00\t4.001095e+00\t1.000000\t0.546000\t162.558465\t4.878535\t\n",
            "000500\t(0500.0)\t6.609384e-03\t3.858012e+00\t3.892105e+00\t1.000000\t0.607000\t240.651704\t4.887955\t\n",
            "000750\t(0750.0)\t6.405959e-03\t3.734019e+00\t3.769034e+00\t1.000000\t0.620000\t331.370819\t5.294681\t\n",
            "000753\t(0753.0)\t6.403055e-03\t3.732584e+00\t3.767607e+00\t1.000000\t0.622000\t339.690262\t5.295148\t\n",
            "001000\t(1000.0)\t6.128212e-03\t3.618819e+00\t3.653923e+00\t1.000000\t0.626000\t428.619216\t5.293156\t\n",
            "001250\t(1250.0)\t5.782982e-03\t3.512686e+00\t3.547766e+00\t1.000000\t0.630000\t519.163821\t4.881382\t\n",
            "001500\t(1500.0)\t5.378769e-03\t3.416091e+00\t3.451097e+00\t1.000000\t0.633000\t610.928892\t4.883432\t\n",
            "001750\t(1750.0)\t4.925528e-03\t3.329300e+00\t3.364227e+00\t1.000000\t0.637000\t702.328579\t5.294870\t\n",
            "001940\t(1940.0)\t4.555102e-03\t3.269953e+00\t3.304930e+00\t1.000000\t0.638000\t775.041556\t4.888358\t\n",
            "002000\t(2000.0)\t4.434416e-03\t3.252397e+00\t3.287360e+00\t1.000000\t0.636000\t801.262459\t4.877840\t\n",
            "002250\t(2250.0)\t3.917529e-03\t3.185291e+00\t3.220258e+00\t1.000000\t0.638000\t893.076906\t4.958337\t\n",
            "002500\t(2500.0)\t3.387592e-03\t3.127757e+00\t3.162734e+00\t1.000000\t0.642000\t983.312528\t5.299324\t\n",
            "002750\t(2750.0)\t2.857656e-03\t3.079411e+00\t3.114452e+00\t1.000000\t0.643000\t1073.285161\t5.223449\t\n",
            "003000\t(3000.0)\t2.340769e-03\t3.039760e+00\t3.074824e+00\t1.000000\t0.643000\t1164.813317\t4.878101\t\n",
            "003250\t(3250.0)\t1.849657e-03\t3.008190e+00\t3.043247e+00\t1.000000\t0.644000\t1257.180832\t5.294868\t\n",
            "003500\t(3500.0)\t1.396416e-03\t2.983985e+00\t3.019088e+00\t1.000000\t0.643000\t1351.223951\t5.298793\t\n",
            "003750\t(3750.0)\t9.922029e-04\t2.966324e+00\t3.001445e+00\t1.000000\t0.643000\t1444.194509\t4.889918\t\n",
            "004000\t(4000.0)\t6.469724e-04\t2.954296e+00\t2.989431e+00\t1.000000\t0.640000\t1535.169996\t5.298050\t\n",
            "004250\t(4250.0)\t3.692257e-04\t2.946901e+00\t2.982043e+00\t1.000000\t0.640000\t1625.166799\t5.296432\t\n",
            "004500\t(4500.0)\t1.658007e-04\t2.943062e+00\t2.978206e+00\t1.000000\t0.642000\t1715.971992\t4.878840\t\n",
            "004750\t(4750.0)\t4.170681e-05\t2.941636e+00\t2.976783e+00\t1.000000\t0.641000\t1807.775999\t4.885040\t\n",
            "004999\t(4999.0)\t6.057487e-10\t2.941434e+00\t2.976582e+00\t1.000000\t0.641000\t1898.692093\t5.302351\t\n",
            "005000\t(5000.0)\t0.000000e+00\t2.941434e+00\t2.976582e+00\t1.000000\t0.641000\t1906.318331\t5.297362\t\n",
            "Training finished\n",
            "Step\tEpoch\tLearning Rate\tTrain Loss\tValid Loss\tTest Loss\tTrain Acc\tValid Acc\tTest Acc\tTime Elapsed\n",
            "005000\t(5000.0)\t0.000000e+00\t2.941434e+00\t2.976582e+00\t2.977538e+00\t1.000000\t0.641000\t0.631400\t1967.782253\t\n"
          ]
        }
      ],
      "source": [
        "tf.config.experimental.set_visible_devices([], 'GPU')\n",
        "\n",
        "config = get_config()\n",
        "data = load_kip_data(config)\n",
        "measurements = run_trainer(data, config)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "BwO9GQYkmWfI"
      ],
      "name": "nn_training.ipynb",
      "provenance": [
        {
          "file_id": "1nlmvs7zUsNLWALub9fr8OBNcAmRYqXib",
          "timestamp": 1631839500786
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
